<#
.SYNOPSIS
WSL 集群部署工具 - 支持批量创建预配置节点

.DESCRIPTION
基于指定镜像创建多个 WSL 实例，并自动完成 SSH 及用户配置

.EXAMPLE
.\Create-WSLCluster.ps1 -BaseImage "Ubuntu-22.04" -Location "F:\wsl" -NodeCount 3
#>

param (
    [Parameter()]
    [string]$BaseImage = "Ubuntu-22.04",

    [Parameter()]
    [string]$Location = "F:\WSL",

    [Parameter()]
    [string]$NamePrefix = "node",

    [Parameter()]
    [int]$NodeCount = 3
)

# 生成节点列表
$nodes = 1..$NodeCount | ForEach-Object {
    @{
        Name = "$NamePrefix$_"
        Path = "$Location\$NamePrefix$_"
    }
}

# 创建基础镜像模板
$baseDistroName = "BaseTemplate"
if (-not (wsl -l | Select-String $baseDistroName)) {
    Write-Host "正在创建基础镜像..."
    wsl --install -d $BaseImage --name $baseDistroName
    wsl --terminate $baseDistroName
}

# 导出基础镜像
$baseVhdPath = "$Location\base.tar"
if (-not (Test-Path $baseVhdPath)) {
    wsl --export $baseDistroName $baseVhdPath
    if (-not $?) {
        Write-Error "Failed to export base image"
        exit 1
    }
}

# 批量创建节点
# Fix base image creation
$baseVhdPath = "$Location\base.tar"
if (-not (Test-Path $baseVhdPath)) {
    Write-Host "Creating base image..."
    wsl --export $BaseImage $baseVhdPath
    if (-not $?) {
        Write-Error "Failed to export base image"
        exit 1
    }
}

# Fix node creation with proper import
$nodes | ForEach-Object {
    $node = $_
    if (-not (Test-Path $node.Path)) {
        New-Item -Path $node.Path -ItemType Directory -Force | Out-Null
    }
    
    Write-Host "Creating node: $($node.Name)"
    wsl --import $node.Name $node.Path $baseVhdPath --version 2
    wsl --set-default-version 2
}

# Fix configuration execution
$nodes | ForEach-Object {
    $nodeName = $_.Name
    Write-Host "Configuring: $nodeName"
    
    # Modified configuration script
    $configScript = @"
    # Base configuration
    apt update && apt upgrade -y
    ldconfig 2>/dev/null || true  # Suppress libcuda link warning
    
    # Rest of the configuration remains unchanged
    apt install -y openssh-server
    useradd -m -s /bin/bash wsluser
    mkdir -p /home/wsluser/.ssh
    echo '$publicKey' > /home/wsluser/.ssh/authorized_keys
    chown -R wsluser:wsluser /home/wsluser/.ssh
    chmod 700 /home/wsluser/.ssh
    chmod 600 /home/wsluser/.ssh/authorized_keys
    
    # SSH configuration
    sed -i 's/#PasswordAuthentication yes/PasswordAuthentication no/' /etc/ssh/sshd_config
    /usr/sbin/sshd -t && /usr/sbin/sshd -D &  # 尝试重新启动 SSH 服务
"@

    wsl -d $nodeName -u root -e bash -c $configScript
}

# 生成 SSH 密钥对
# 确保 $Location 有值
if ([string]::IsNullOrWhiteSpace($Location)) {
    Write-Error "Location 参数为空，请提供有效的路径。"
    exit 1
}
$sshKeyPath = Join-Path -Path $Location -ChildPath "cluster_key"
if (-not (Test-Path -Path $sshKeyPath)) {
    ssh-keygen -t rsa -b 4096 -f $sshKeyPath -N "" -q
}
$publicKey = Get-Content -Path "$sshKeyPath.pub"

# 节点统一配置
$nodes | ForEach-Object {
    $nodeName = $_.Name
    Write-Host "配置节点: $nodeName"
    
    # 基础配置
    $configScript = @"
    # 用户及权限
    sudo adduser --disabled-password --gecos "" wsluser
    echo 'wsluser ALL=(ALL) NOPASSWD:ALL' | sudo tee /etc/sudoers.d/wsluser
    sudo mkdir -p /home/wsluser/.ssh
    echo '$publicKey' | sudo tee /home/wsluser/.ssh/authorized_keys
    sudo chown -R wsluser:wsluser /home/wsluser/.ssh
    sudo chmod 700 /home/wsluser/.ssh
    sudo chmod 600 /home/wsluser/.ssh/authorized_keys

    # SSH 配置
    sudo sed -i 's/#PasswordAuthentication yes/PasswordAuthentication no/' /etc/ssh/sshd_config
    sudo /usr/sbin/sshd -t && sudo /usr/sbin/sshd -D &  # 尝试重新启动 SSH 服务
"@

    wsl -d $nodeName -- bash -c $configScript
}

# 输出节点信息
Write-Host "`nCluster nodes:"
$nodes | ForEach-Object {
    Write-Host "- $($_.Name)"
}
Write-Host "`n使用以下密钥连接节点:"
Write-Host "SSH 私钥路径: $sshKeyPath"
# 确保中文显示正常，这里重新修正引号
Write-Host "SSH 私钥路径: $sshKeyPath"
